(defpackage :tree
  (:use :cl)
  (:export base-node
	   node
	   multi-parent-node
	   parent
	   parents
	   children
	   delete-children
	   child
	   single-child-mixin
           copy-node
           dumb-copy
           node-equal
           replace-child
           recreate-parent-slots
           sanity-check
           descendant-p
           common-ancestor
           paths-by
           paths-by-parents
           paths-by-children
           with-iterator
           recurse
           depth
           detect-loop))

(in-package :tree)

(defclass base-node ()
  ((parents :initform nil :initarg :parents :accessor parents)
   (children :initform nil :initarg :children :accessor children)))

(defclass node (base-node)
  ()
  (:documentation "A node that has only one parent."))

(defgeneric parent (node)
  (:documentation "For single parent nodes, return that single parent."))
(defgeneric (setf parent) (value node)
  (:documentation "For single parent nodes, return that single parent."))

(defmethod parent ((node node))
  (assert (<= (length (parents node)) 1))
  (first (parents node)))

(defmethod (setf parent) ((parent node) (obj node))
  (setf (parents obj) (list parent)))

(defclass multi-parent-node (base-node)
  ()
  (:documentation "A node that can have more than 1 parent."))

(defmethod shared-initialize :after ((obj base-node) slots &key parent)
  (when parent
    (setf (slot-value obj 'parents) (list parent)))
  (dolist (i (children obj))
    (pushnew obj (slot-value i 'parents)))
  (dolist (i (parents obj))
    (pushnew obj (slot-value i 'children))))

(defmethod (setf parents) :before (parents (obj node))
  (dolist (i (parents obj))
    (setf (slot-value i 'children) (remove obj (slot-value i 'children)))))

(defmethod (setf parents) :after (parents (obj node))
  (dolist (i (parents obj))
    (pushnew obj (slot-value i 'children))))

(defmethod (setf children) :before (children (obj node))
;;   (dolist (i (children obj))
;;     (setf (slot-value i 'parents) nil))
  )

(defmethod (setf children) :after (children (obj node))
  (dolist (i children)
    (push obj (slot-value i 'parents))))

(defgeneric delete-children (obj &rest children))

(defmethod delete-children ((obj node) &rest children)
  (setf (children obj) (set-difference (children obj) children)))

(defclass single-child-mixin (base-node)
  ())

(defmethod shared-initialize :after ((obj single-child-mixin) slots &key child)
  (unless (children obj)
    (setf (children obj) (when child
                           (list child)))))

(defgeneric child (obj))
(defgeneric (setf child) (child obj))

(defmethod (setf child) ((child node) (obj single-child-mixin))
  (setf (children obj) (list child)))

(defmethod (setf child) ((child null) (obj single-child-mixin))
  (setf (children obj) nil))

;; FIXME: this should only be applicable to single-child-mixin?
(defmethod child ((obj node))
  (first (children obj)))

(defun descendant-p (node ancestor)
  "Return T if node is a descendant of ascestor."
  (or (find node (children ancestor))
      (some (lambda (n) (descendant-p node n)) (children ancestor))))

(defun common-ancestor (a b)
  "Return the closest common ancestor that a and b have."
  (labels ((common (x node)
             (if (descendant-p x node)
                 node
                 (loop for i in (tree:parents node)
                      for c = (common x i)
                      when c return c))))
    (common a b)))

(defun paths-by (start destination via-fn)
  "Return all the paths that go from start to destination by way of calling VIA-FN."
  (labels ((paths (current)
             (if (eq current destination)
                 `((,destination))
                 (let ((nodes (funcall via-fn current)))
                   (delete nil
                           (mapcan (lambda (c)
                                     (loop for i in (paths c)
                                        when i
                                        collect (cons current i)))
                                   nodes))))))
    (paths start)))

(defun paths-by-parents (start destination)
  (paths-by start destination 'parents))

(defun paths-by-children (start destination)
  (paths-by start destination 'children))

(defstruct multi-parent-collector
  alist)

(defmethod copy-node ((node base-node) deep &optional (mp (make-multi-parent-collector)))
  (labels ((copy ()
             (make-instance (type-of node)
                            :children (when deep
                                        (mapcar (lambda (c) (copy-node c deep mp)) (children node))))))
    (if (and deep
             (> (length (parents node)) 1))
        (let ((hit (assoc node (multi-parent-collector-alist mp))))
          (if hit
              (cdr hit)
              (let ((new (copy)))
                (u:copy-slots node new)
                (setf (multi-parent-collector-alist mp)
                      (acons node new (multi-parent-collector-alist mp)))
                new)))
        (let ((new (copy)))
          (u:copy-slots node new)
          new))))

(defmethod dumb-copy ((node base-node) deep)
  "Don't handle multi-parent nodes specially."
  (labels ((copy ()
             (make-instance (type-of node)
                            :children (when deep
                                        (mapcar (lambda (c) (dumb-copy c deep)) (children node))))))
    (let ((new (copy)))
      (u:copy-slots node new)
      new)))

;; (defun copy-node (tree-node deep)
;;   (let ((mp nil))
;;     (labels ((lookup (node)
;;                (when (and deep
;;                           (> (length (parents node)) 1))
;;                  (let ((hit (assoc node mp)))
;;                    (when hit
;;                      (cdr hit)))))
;;              (memoize (node new)
;;                (when (> (length (parents node)) 1)
;;                  (setf mp
;;                        (acons node new mp))))
;;              (copy (node)
;;                (let ((c (make-instance (type-of node))))
;;                  (u:copy-slots node c)
;;                  c))
;;              (iterate (node)
;;                (let ((new (or (lookup node)
;;                               (copy node))))
;;                  (when deep
;;                    (memoize node new)
;;                    (setf (children new)
;;                          (mapcar #'iterate (children node))))
;;                  new)))
;;       (iterate tree-node))))

(defgeneric node-equal (node1 node2)
  (:method-combination and))

(defmethod node-equal and ((node1 base-node) (node2 base-node))
  (and
   (eq (type-of node1) (type-of node2))
   (eql (length (children node1))
        (length (children node2)))
   (loop
      for c1 in (children node1)
      for c2 in (children node2)
      always (node-equal c1 c2))))

(defun replace-child (parent old new)
  (setf (slot-value parent 'children) (substitute new old (children parent)))
  (pushnew parent (slot-value new 'parents)))

(defun recreate-parent-slots (tree)
  "Rebuild the parent slots in TREE's nodes and return tree."
  (labels ((clear (node)
             (setf (slot-value node 'parents) nil)
             (mapc #'clear (children node)))
           (set (node parent)
             (when parent
               (pushnew parent (slot-value node 'parents)))
             (dolist (i (children node))
               (set i node))))
    (clear tree)
    (set tree nil)
    tree))

(defun sanity-check (tree)
  "Check that a node's parents all have the node as a child and that
all parents are a descendant of TREE. Return all nodes that do not
comply. Used for debugging."
  (labels ((check (node all)
             (append
              (unless (and (or (null (children node))
                               (every (lambda (c) (find node (parents c))) (children node)))
                           (every (lambda (p) (find p all)) (parents node)))
                (list node))
              (mapcan (lambda (c) (check c all)) (children node))))
           (nodes (node)
             (cons node (mapcan #'nodes (children node)))))
    (check tree (nodes tree))))

(defmacro with-iterator ((iterate-name node &optional (recurse 'recurse) (type t)) iterate-statement &body body)
  (assert (u:neq iterate-name recurse)
          ()
          "iterate-name cannot have the same value as recurse")
  (u:with-gensyms (mp)
    `(let (,mp)
       (labels ((,iterate-name (,node)
                  (labels ((,recurse () (mapc #',iterate-name (children ,node))))
                    (if (typep ,node ,type)
                        (when (or (<= (length (parents ,node)) 1)
                                  (not (find ,node ,mp)))
                          (when (> (length (parents ,node)) 1)
                            (push ,node ,mp))
                          ,iterate-statement)
                        (,recurse)))))
         (progn ,@body)))))

(defun depth (node &optional (deepest t))
  "Return the depth of the node in its tree. If multiple parents are
encountered and MAX is T then return the deepest path. Otherwise
return the shallowest."
  (labels ((d (node depth)
             (if (parents node)
                 (apply (if deepest 'max 'min) (mapcar (lambda (n) (d n (1+ depth))) (parents node)))
                 depth)))
    (d node 0)))

(defun detect-loop (tree)
  (labels ((walk (node seen)
             (if (find node seen)
                 t
                 (loop for n in (tree:children node)
                      thereis (walk n (cons node seen))))))
    (walk tree nil)))